"""
知识库服务模块
处理RAG知识库的加载、管理和检索功能
"""

import logging
from pathlib import Path
from typing import List, Optional

from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.document_loaders import TextLoader
from langchain_community.embeddings import DashScopeEmbeddings
from langchain_community.vectorstores import FAISS
from langchain_core.documents import Document

from backend.config.constants import Config

logger = logging.getLogger(__name__)


class KnowledgeService:
    """知识库服务类"""
    
    def __init__(self):
        self.knowledge_base_path = Path("./knowledge")
        self.vector_store: Optional[FAISS] = None
        self.embeddings = None
        self._initialize_embeddings()
    
    def _initialize_embeddings(self):
        """初始化DashScope嵌入模型"""
        try:
            if Config.DASHSCOPE_API_KEY:
                self.embeddings = DashScopeEmbeddings(
                    dashscope_api_key=Config.DASHSCOPE_API_KEY,
                    model=Config.EMBEDDING_MODEL
                )
                logger.info("成功初始化DashScope嵌入模型")
            else:
                logger.error("未设置DASHSCOPE_API_KEY环境变量")
                raise ValueError("请设置DASHSCOPE_API_KEY环境变量")
        except Exception as e:
            logger.error(f"初始化DashScope嵌入模型失败: {e}")
            raise
    
    def load_knowledge_files(self) -> List[Document]:
        """加载知识库文件"""
        documents = []
        
        if not self.knowledge_base_path.exists():
            logger.warning(f"知识库目录不存在: {self.knowledge_base_path}")
            return documents
        
        # 遍历知识库目录中的所有txt文件
        for file_path in self.knowledge_base_path.glob("*.txt"):
            try:
                logger.info(f"加载知识库文件: {file_path}")
                loader = TextLoader(str(file_path), encoding="utf-8")
                docs = loader.load()
                documents.extend(docs)
            except Exception as e:
                logger.error(f"加载文件 {file_path} 失败: {e}")
                continue
        
        logger.info(f"共加载 {len(documents)} 个文档")
        return documents
    
    def create_vector_store(self) -> Optional[FAISS]:
        """创建向量存储"""
        try:
            # 加载知识库文件
            documents = self.load_knowledge_files()
            
            if not documents:
                logger.warning("没有加载到任何知识库文档")
                return None
            
            # 分割文档
            text_splitter = RecursiveCharacterTextSplitter(
                chunk_size=1000,
                chunk_overlap=200,
                length_function=len,
            )
            splits = text_splitter.split_documents(documents)
            
            logger.info(f"文档分割为 {len(splits)} 个片段")
            
            # 创建向量存储
            if self.embeddings:
                self.vector_store = FAISS.from_documents(splits, self.embeddings)
                logger.info("成功创建向量存储")
                return self.vector_store
            else:
                logger.error("嵌入模型未初始化，无法创建向量存储")
                return None
                
        except Exception as e:
            logger.error(f"创建向量存储失败: {e}")
            return None
    
    def search_knowledge(self, query: str, k: int = 4) -> List[Document]:
        """检索相关知识"""
        if not self.vector_store:
            logger.warning("向量存储未初始化")
            return []
        
        try:
            # 搜索相关文档
            docs = self.vector_store.similarity_search(query, k=k)
            logger.info(f"检索到 {len(docs)} 个相关文档")
            return docs
        except Exception as e:
            logger.error(f"知识检索失败: {e}")
            return []
    
    def get_retriever(self):
        """获取检索器"""
        if self.vector_store:
            return self.vector_store.as_retriever()
        return None
    
    def is_initialized(self) -> bool:
        """检查服务是否已初始化"""
        return self.vector_store is not None


# 全局知识库服务实例
knowledge_service = KnowledgeService()